Phase Mask Design¶

In this notebook, we will illustrate the problem of inverse design of a phase mask: we will choose the example from Wong et al, 2021, designing a diffractive pupil phase mask for the Toliman telescope.

In order to get high precision centroids, we need to maximize the gradient energy of the pupil; in order to satisfy fabrication constraints, we need a binary mask with phases only in {0, π}.

In [1]:

import jax
import jax.numpy as np
from jax import vmap

import matplotlib.pyplot as plt
from tqdm import tqdm

%matplotlib inline
plt.rcParams['image.cmap'] = 'hot'
plt.rcParams["text.usetex"] = 'false'
plt.rcParams['figure.dpi'] = 120

import jax import jax.numpy as np from jax import vmap import matplotlib.pyplot as plt from tqdm import tqdm %matplotlib inline plt.rcParams['image.cmap'] = 'hot' plt.rcParams["text.usetex"] = 'false' plt.rcParams['figure.dpi'] = 120

/Users/louis/mambaforge/envs/dlux/lib/python3.10/site-packages/jax/_src/lib/__init__.py:34: UserWarning: JAX on Mac ARM machines is experimental and minimally tested. Please see https://github.com/google/jax/issues/5501 in the event of problems.
  warnings.warn("JAX on Mac ARM machines is experimental and minimally tested. "

We will first generate an orthonormal basis for the pupil phases, and then threshold this to {0, 1} while preserving soft edges using the Continuous Latent Image Mask Binarization (CLIMB) algorithm from the Wong et al paper.

In [2]:

from sklearn.decomposition import PCA

from sklearn.decomposition import PCA

Generate the support of the pupil:

In [3]:

wf_npix = 256
oversample = 3
nslice = 3

npix = wf_npix * oversample
c = (npix - 1) / 2.
xs = (np.arange(npix) - c) / c
XX, YY = np.meshgrid(xs, xs)
RR = np.sqrt(XX ** 2 + YY ** 2)
PHI = np.arctan2(YY, XX)

mask = np.logical_and(RR <= 1, RR >= 0.175).astype(float)
plt.imshow(mask)
plt.colorbar()
plt.show()

wf_npix = 256 oversample = 3 nslice = 3 npix = wf_npix * oversample c = (npix - 1) / 2. xs = (np.arange(npix) - c) / c XX, YY = np.meshgrid(xs, xs) RR = np.sqrt(XX ** 2 + YY ** 2) PHI = np.arctan2(YY, XX) mask = np.logical_and(RR <= 1, RR >= 0.175).astype(float) plt.imshow(mask) plt.colorbar() plt.show()

WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

Generate basis vectors however you like - in this case we are using logarithmic radial harmonics and sines and cosines in θ., but you can do whatever you like here. This code is not important; just generate your favourite not-necessarily-orthonormal basis, and we will use PCA to orthonormalize it later on.

In [4]:

a = 10
b = 8
ith = 10

As = np.arange(-a, a+1)
Bs = 3 * np.arange(0, b+1)
Cs = np.array([-np.pi/2, np.pi/2])
Is = np.arange(-ith, ith+1)

LRHF_fn = lambda A, B, C, RR, PHI: np.cos(A*np.log(RR + 1e-12) + B*PHI + C)
sine_fn = lambda i, RR: np.sin(i * np.pi * RR)
cose_fn = lambda i, RR: np.cos(i * np.pi * RR)

gen_LRHF_basis = vmap(vmap(vmap(LRHF_fn, (None, 0, None, None, None)), (0, None, None, None, None)), (None, None, 0, None, None))
gen_sine_basis = vmap(sine_fn, in_axes=(0, None))
gen_cose_basis = vmap(cose_fn, in_axes=(0, None))

LRHF_basis = gen_LRHF_basis(As, Bs, Cs, RR, PHI).reshape([len(As)*len(Bs)*len(Cs), npix, npix])
sine_basis = gen_sine_basis(Is, RR)
cose_basis = gen_cose_basis(Is, RR)

LRHF_flat = LRHF_basis.reshape([len(As)*len(Bs)*len(Cs), npix*npix])
sine_flat = sine_basis.reshape([len(sine_basis), npix*npix])
cose_flat = cose_basis.reshape([len(cose_basis), npix*npix])

full_basis = np.concatenate([
    LRHF_flat,
    sine_flat,
    cose_flat
])

a = 10 b = 8 ith = 10 As = np.arange(-a, a+1) Bs = 3 * np.arange(0, b+1) Cs = np.array([-np.pi/2, np.pi/2]) Is = np.arange(-ith, ith+1) LRHF_fn = lambda A, B, C, RR, PHI: np.cos(A*np.log(RR + 1e-12) + B*PHI + C) sine_fn = lambda i, RR: np.sin(i * np.pi * RR) cose_fn = lambda i, RR: np.cos(i * np.pi * RR) gen_LRHF_basis = vmap(vmap(vmap(LRHF_fn, (None, 0, None, None, None)), (0, None, None, None, None)), (None, None, 0, None, None)) gen_sine_basis = vmap(sine_fn, in_axes=(0, None)) gen_cose_basis = vmap(cose_fn, in_axes=(0, None)) LRHF_basis = gen_LRHF_basis(As, Bs, Cs, RR, PHI).reshape([len(As)*len(Bs)*len(Cs), npix, npix]) sine_basis = gen_sine_basis(Is, RR) cose_basis = gen_cose_basis(Is, RR) LRHF_flat = LRHF_basis.reshape([len(As)*len(Bs)*len(Cs), npix*npix]) sine_flat = sine_basis.reshape([len(sine_basis), npix*npix]) cose_flat = cose_basis.reshape([len(cose_basis), npix*npix]) full_basis = np.concatenate([ LRHF_flat, sine_flat, cose_flat ])

Orthonormalize with PCA - could also use Gram-Schmidt if you prefer.

In [5]:

%%time
pca = PCA().fit(full_basis)

components = pca.components_.reshape([len(full_basis), npix, npix])
components = np.copy(components[:99,:,:])
basis = np.concatenate([np.mean(pca.mean_)*np.array(np.ones((1,npix,npix))), components])

%%time pca = PCA().fit(full_basis) components = pca.components_.reshape([len(full_basis), npix, npix]) components = np.copy(components[:99,:,:]) basis = np.concatenate([np.mean(pca.mean_)*np.array(np.ones((1,npix,npix))), components])

CPU times: user 2min 36s, sys: 1.19 s, total: 2min 37s
Wall time: 29.5 s

Show the pretty basis vectors:

In [6]:

nfigs = 100

ncols = 10
nrows = 1 + nfigs//ncols

plt.figure(figsize=(4*ncols, 4*nrows))
for i in range(nfigs):
    plt.subplot(nrows, ncols, i+1)
    plt.imshow(basis[i], cmap='seismic')
    plt.xticks([])
    plt.yticks([])
plt.tight_layout()

nfigs = 100 ncols = 10 nrows = 1 + nfigs//ncols plt.figure(figsize=(4*ncols, 4*nrows)) for i in range(nfigs): plt.subplot(nrows, ncols, i+1) plt.imshow(basis[i], cmap='seismic') plt.xticks([]) plt.yticks([]) plt.tight_layout()

In [7]:

# save for use later
np.save('files/basis', basis)

# save for use later np.save('files/basis', basis)

Optimizing the Pupil¶

First, we will import some more functions from dLux:

In [8]:

import equinox as eqx # object-oriented jax
import optax # optimization
import functools as ft # used for equinox filtering

import dLux as dl

from dLux.utils.models import toliman_layers # pre-built Toliman pupils
from dLux.utils.plotting import plot_batch
from dLux.utils.helpers import get_ppf, norm_rand # pupil scaling

import equinox as eqx # object-oriented jax import optax # optimization import functools as ft # used for equinox filtering import dLux as dl from dLux.utils.models import toliman_layers # pre-built Toliman pupils from dLux.utils.plotting import plot_batch from dLux.utils.helpers import get_ppf, norm_rand # pupil scaling

Initialize a wavelength grid and decide on the mode coordinate grid scaling relative to the pupil:

In [9]:

wavels = 1e-9 * np.linspace(595, 695, 3) # wavelengths always in m

# rescale 
ppf = get_ppf(wavels)
min_fringe = 1
max_fringe = 10

rmin = min_fringe/ppf
rmax = max_fringe/ppf
print(rmin, rmax)

wavels = 1e-9 * np.linspace(595, 695, 3) # wavelengths always in m # rescale ppf = get_ppf(wavels) min_fringe = 1 max_fringe = 10 rmin = min_fringe/ppf rmax = max_fringe/ppf print(rmin, rmax)

5.2393847 52.39385

Now we initialize an optical system using the default Toliman model, plus the CLIMB layer that soft-thresholds the edges (see Wong et al, 2021, sec 3.2.2). In brief, what this does is make an OPD map as a weighted sum of modes; where this is positive, we set to π phase, where it is negative, we set it to zero, and on the edges, we soft-edge it to propagate gradients.

In [10]:

coeffs = norm_rand(0, 100, len(basis))

layers = toliman_layers([dl.ApplyBasisCLIMB(basis, coeffs, np.mean(wavels))])    
osys = dl.OpticalSystem(layers, wavels=wavels, detector_layers=[dl.ApplyJitter(1)])

coeffs = norm_rand(0, 100, len(basis)) layers = toliman_layers([dl.ApplyBasisCLIMB(basis, coeffs, np.mean(wavels))]) osys = dl.OpticalSystem(layers, wavels=wavels, detector_layers=[dl.ApplyJitter(1)])

Visualize this:

In [11]:

psf = osys()

fig, (ax1,ax2) = plt.subplots(1,2,figsize=(6.0,3.3))
ax1.imshow(osys.layers[2].array*osys.layers[4].get_binary_phase(),cmap='hot')
ax1.set_title('Input OPD')

psf_in = osys()
ax2.imshow(psf_in)
ax2.set_title('Input PSF')

psf = osys() fig, (ax1,ax2) = plt.subplots(1,2,figsize=(6.0,3.3)) ax1.imshow(osys.layers[2].array*osys.layers[4].get_binary_phase(),cmap='hot') ax1.set_title('Input OPD') psf_in = osys() ax2.imshow(psf_in) ax2.set_title('Input PSF')

Out[11]:

Text(0.5, 1.0, 'Input PSF')

Now we want to do some optimization! Using the filter_spec syntax, we allow it to optimize only the coefficients in the CLIMB basis.

In [12]:

# Default values to group 0
filter_spec = jax.tree_map(lambda _: False, osys)

# Set filtering Tree
filter_spec = eqx.tree_at(lambda osys: osys.layers[4].coeffs, filter_spec, replace=True)

@eqx.filter_jit
@ft.partial(eqx.filter_value_and_grad, filter_spec=filter_spec)
def loss_func(tel, rmin=20, rmax=150, power=0.5):
    # Get PSF, Gradients and Mask
    psf = tel()
    Rmask = dl.utils.get_Rmask(psf.shape[0], rmin, rmax)

    # Calcualte loss
    loss1 = - np.power(Rmask*dl.utils.get_GE(psf), power).sum()
    loss2 = - np.power(Rmask*dl.utils.get_RGE(psf), power).sum()
    loss = loss1 + loss2
    return loss

# Default values to group 0 filter_spec = jax.tree_map(lambda _: False, osys) # Set filtering Tree filter_spec = eqx.tree_at(lambda osys: osys.layers[4].coeffs, filter_spec, replace=True) @eqx.filter_jit @ft.partial(eqx.filter_value_and_grad, filter_spec=filter_spec) def loss_func(tel, rmin=20, rmax=150, power=0.5): # Get PSF, Gradients and Mask psf = tel() Rmask = dl.utils.get_Rmask(psf.shape[0], rmin, rmax) # Calcualte loss loss1 = - np.power(Rmask*dl.utils.get_GE(psf), power).sum() loss2 = - np.power(Rmask*dl.utils.get_RGE(psf), power).sum() loss = loss1 + loss2 return loss

/Users/louis/mambaforge/envs/dlux/lib/python3.10/site-packages/equinox/grad.py:71: UserWarning: For brevity the `filter_spec` argument has been renamed `arg`
  warnings.warn("For brevity the `filter_spec` argument has been renamed `arg`")

Evaluate once to jit compile:

In [13]:

%%time
loss, grads = loss_func(osys, rmin=rmin, rmax=rmax) # Compile

print("Initial Loss: {}".format(loss))

%%time loss, grads = loss_func(osys, rmin=rmin, rmax=rmax) # Compile print("Initial Loss: {}".format(loss))

Initial Loss: -32.77363586425781
CPU times: user 4 s, sys: 78.3 ms, total: 4.07 s
Wall time: 2.78 s

Gradient descent time!

In [14]:

model = osys
optim = optax.adam(2.5e1)
opt_state = optim.init(model)

errors, grads_out, models_out = [], [], []

with tqdm(range(100),desc='Gradient Descent') as t:
    for i in t: 
        loss, grads = loss_func(model, rmin=rmin, rmax=rmax)
        updates, opt_state = optim.update(grads, opt_state)
        model = eqx.apply_updates(model, updates)
        models_out.append(model)
        errors.append(loss)
        grads_out.append(grads)

        t.set_description("Loss: {:.3f}".format(loss)) # update the progress bar

model = osys optim = optax.adam(2.5e1) opt_state = optim.init(model) errors, grads_out, models_out = [], [], [] with tqdm(range(100),desc='Gradient Descent') as t: for i in t: loss, grads = loss_func(model, rmin=rmin, rmax=rmax) updates, opt_state = optim.update(grads, opt_state) model = eqx.apply_updates(model, updates) models_out.append(model) errors.append(loss) grads_out.append(grads) t.set_description("Loss: {:.3f}".format(loss)) # update the progress bar

Loss: -47.506: 100%|█████████████████████████████████████████████████████████████████| 100/100 [00:24<00:00,  4.13it/s]

Visualize the results:

In [15]:

coeffs_out = np.array([model.layers[4].coeffs for model in models_out])

mosaic = """
    AABB
    CDEF
    """

fig = plt.figure(constrained_layout=True,figsize=(12, 8))
axes = fig.subplot_mosaic(mosaic)

for ax in ['C','D','E','F']:
    axes[ax].set_xticks([])
    axes[ax].set_yticks([])


axes['A'].plot(np.array(errors))
axes['A'].set_title("Loss")

axes['B'].set_title("Coeffs")
axes['B'].plot(coeffs_out[:], alpha=0.2)

model_in = osys
axes['C'].imshow(model_in.layers[2].array*model_in.layers[4].get_binary_phase(),cmap='hot')
axes['C'].set_title('Input OPD')

psf_in = model_in()
axes['D'].imshow(psf_in)
axes['D'].set_title('Input PSF')

final = models_out[-1]
axes['E'].imshow(final.layers[2].array*final.layers[4].get_binary_phase(),cmap='hot')
axes['E'].set_title('Final OPD')

psf_out = final()
axes['F'].imshow(psf_out)
axes['F'].set_title('Final PSF')

coeffs_out = np.array([model.layers[4].coeffs for model in models_out]) mosaic = """ AABB CDEF """ fig = plt.figure(constrained_layout=True,figsize=(12, 8)) axes = fig.subplot_mosaic(mosaic) for ax in ['C','D','E','F']: axes[ax].set_xticks([]) axes[ax].set_yticks([]) axes['A'].plot(np.array(errors)) axes['A'].set_title("Loss") axes['B'].set_title("Coeffs") axes['B'].plot(coeffs_out[:], alpha=0.2) model_in = osys axes['C'].imshow(model_in.layers[2].array*model_in.layers[4].get_binary_phase(),cmap='hot') axes['C'].set_title('Input OPD') psf_in = model_in() axes['D'].imshow(psf_in) axes['D'].set_title('Input PSF') final = models_out[-1] axes['E'].imshow(final.layers[2].array*final.layers[4].get_binary_phase(),cmap='hot') axes['E'].set_title('Final OPD') psf_out = final() axes['F'].imshow(psf_out) axes['F'].set_title('Final PSF')

Out[15]:

Text(0.5, 1.0, 'Final PSF')

In [16]:

# Save mask for use in flatfield_calibration notebook
mask_out = models_out[-1].layers[4].get_binary_phase()
np.save("files/test_mask", mask_out)

# Save mask for use in flatfield_calibration notebook mask_out = models_out[-1].layers[4].get_binary_phase() np.save("files/test_mask", mask_out)

In [ ]: